"""
将地震目录中的地震事件聚类
"""
from turtle import color
import pandas as pd
import numpy as np
from sklearn.cluster import KMeans
from obspy.core.utcdatetime import UTCDateTime
from constant import CATALOG_PATH, OK027_POS, OK029_POS, EVENT_CLUSTER, \
    FILTER_LAT, FILTER_LON, CATALOG_CLUSTER_PATH
from collections import Counter
import matplotlib.pyplot as plt
import brewer2mpl


def filter_catalog(cat):
    """ Filter the events in cat
    to keep the events near Guthrie 
    """
    cat = cat[(cat.latitude > FILTER_LAT[0]) & (cat.latitude < FILTER_LAT[1])
              & (cat.longitude > FILTER_LON[0]) & (cat.longitude < FILTER_LON[1])]
    # Filter event after 15th February 2014
    begin_february = UTCDateTime(2014, 2, 15, 00, 00, 00, 500000)
    cat = cat[cat.utc_timestamp > begin_february]
    return cat


def main():
    catalog_path = CATALOG_PATH
    cat = pd.read_csv(catalog_path)
    cat = filter_catalog(cat)
    lat = cat['latitude'].values
    lon = cat['longitude'].values
    # depth = cat['depth'].values
    # utc_time = cat["utc_timestamp"].values

    print(f"Number of events to cluster: {len(lat)}", flush=True)

    feats = np.hstack((lon[:, None], lat[:, None]))

    initialization = np.array([[-97.6, 36],
                               [-97.4, 35.85],
                               [-97.2, 35.85],
                               [-97.3, 35.75],
                               [-97.4, 35.95],
                               [-97.6, 35.75]])
    clust = KMeans(EVENT_CLUSTER, init=initialization, n_init=1).fit(feats)
    cluster_result = clust.labels_
    count = Counter(cluster_result)
    labels = list(set(cluster_result))
    print(count, flush=True)
    cluster_center = clust.cluster_centers_
    # 绘图
    # cluster_show(lon, lat, clust, labels, cluster_result,cluster_center)

    # 保存结果
    res = cluster_result.reshape(-1, 1)
    cat['cluster_id'] = res
    cat['cluster_center'] = [cluster_center[i] for i in cluster_result]

    save_path = CATALOG_CLUSTER_PATH
    cat.to_csv(save_path, index=False)
    print("over.")


def cluster_show(lon, lat, clust, labels, cluster_result, cluster_center):

    x = np.linspace(min(lon), max(lon), 1000)
    y = np.linspace(min(lat), max(lat), 1000)
    X, Y = np.meshgrid(x, y)
    XX = np.array([X.ravel(), Y.ravel()]).T
    Z = clust.predict(XX)
    Z = Z.reshape(X.shape)
    plt.contour(X, Y, Z, colors='k', levels=range(EVENT_CLUSTER))
    plt.xlabel("longitude")
    plt.ylabel("latitude")

    plt.plot(OK029_POS['lon'], OK029_POS['lat'], '*', color='r', markersize=12)
    plt.plot(OK027_POS['lon'], OK027_POS['lat'], '*', color='r', markersize=12)

    for label in labels:
        colors = brewer2mpl.get_map('Set2', 'qualitative', EVENT_CLUSTER).mpl_colors
        plt.scatter(lon[cluster_result == label], lat[cluster_result == label],
                    c=[colors[label]], linewidth=0, label=f'type={label}')
    # plt.savefig("cluster_result.jpg")

    plt.show()


if __name__ == "__main__":
    main()
